package org.theSeed.webFlux.filter;

import cn.hutool.core.lang.UUID;
import cn.hutool.core.util.ObjectUtil;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import org.theSeed.webFlux.sysEnum.FilterOrder;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;

/**
 * 日志过滤器
 */
public class TraceFilter implements WebFilter, Ordered {
    public TraceFilter() {
    }

    public TraceFilter(Integer orderNum) {
        this.orderNum = orderNum;
    }

    public TraceFilter(String headerName) {
        this.headerName = headerName;
    }

    public TraceFilter(Integer orderNum, String headerName) {
        this.orderNum = orderNum;
        this.headerName = headerName;
    }

    protected Integer orderNum = FilterOrder.TRACE.getCode();

    /**
     * 头名称
     */
    protected String headerName = "TRACE_ID";

    @Override
    public int getOrder() {
        return orderNum;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        ServerHttpRequest oldRequest = exchange.getRequest();
        ServerHttpRequest tmp = null;
        String traceId = exchange.getRequest().getHeaders().getFirst(headerName);
        if (ObjectUtil.isEmpty(traceId)){
            traceId = UUID.randomUUID().toString();
            String finalTraceId = traceId;
            tmp = oldRequest.mutate()
                    .headers(headers -> {
                        headers.addAll(oldRequest.getHeaders());
                        // 添加自定义请求头
                        headers.add(headerName, finalTraceId);
                    })
                    .build();
        }else{
            tmp = oldRequest;
        }
        String tmpTraceId = traceId;
        MDC.put(headerName,tmpTraceId);
        return chain.filter(exchange.mutate().request(tmp).build()).contextWrite(ctx -> {
            ctx.put(headerName, tmpTraceId);
            return ctx;
        }).doFinally(signalType -> MDC.remove(headerName));
    }
}
